# -*- coding: utf-8 -*-
"""
A股股票标签同步模块
从akshare获取股票的概念、板块、行业标签信息，并同步到数据库
"""
import logging
import pandas as pd
import akshare as ak
from datetime import datetime
from typing import List, Dict, Any, Optional
from database import DatabaseManager, TagDAO
from config import DATABASE_CONFIG, SYNC_CONFIG

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class StockTagSync:
    """A股股票标签同步器"""
    
    def __init__(self):
        self.db_manager = DatabaseManager(DATABASE_CONFIG)
        self.tag_dao = TagDAO(self.db_manager)
        self.batch_size = SYNC_CONFIG.get('batch_size', 1000)
        self.request_delay = SYNC_CONFIG.get('request_delay', 1.0)
    
    def get_concept_boards(self) -> pd.DataFrame:
        """获取概念板块列表"""
        try:
            logger.info("正在获取概念板块列表...")
            df = ak.stock_board_concept_name_ths()
            
            # 重命名列
            df.rename(columns={
                "板块名称": "name",
                "板块代码": "code"
            }, inplace=True)
            
            # 添加类型
            df['type'] = 'concept'
            df['source'] = 'akshare'
            
            logger.info(f"成功获取 {len(df)} 个概念板块")
            return df
            
        except Exception as e:
            logger.error(f"获取概念板块失败: {e}")
            raise
    
    def get_industry_boards(self) -> pd.DataFrame:
        """获取行业板块列表"""
        try:
            logger.info("正在获取行业板块列表...")
            df = ak.stock_board_industry_name_ths()
            
            # 重命名列
            df.rename(columns={
                "板块名称": "name",
                "板块代码": "code"
            }, inplace=True)
            
            # 添加类型
            df['type'] = 'industry'
            df['source'] = 'akshare'
            
            logger.info(f"成功获取 {len(df)} 个行业板块")
            return df
            
        except Exception as e:
            logger.error(f"获取行业板块失败: {e}")
            raise
    
    def get_sector_boards(self) -> pd.DataFrame:
        """获取地域板块列表"""
        try:
            logger.info("正在获取地域板块列表...")
            # akshare没有直接的地域板块接口，我们创建一个空的数据框
            # 或者可以从其他数据源获取地域信息
            df = pd.DataFrame(columns=['name', 'code', 'type', 'source'])
            
            # 添加一些常见的地域标签
            regions = [
                '北京', '上海', '广东', '深圳', '浙江', '江苏', '山东', '福建',
                '四川', '湖北', '湖南', '河南', '安徽', '河北', '辽宁', '江西',
                '山西', '陕西', '重庆', '天津', '广西', '云南', '贵州', '甘肃',
                '青海', '宁夏', '新疆', '西藏', '内蒙古', '黑龙江', '吉林'
            ]
            
            data = []
            for i, region in enumerate(regions, 1):
                data.append({
                    'name': region,
                    'code': f'AREA{i:03d}',
                    'type': 'sector',
                    'source': 'manual'
                })
            
            df = pd.DataFrame(data)
            logger.info(f"成功获取 {len(df)} 个地域板块")
            return df
            
        except Exception as e:
            logger.error(f"获取地域板块失败: {e}")
            raise
    
    def get_board_stocks(self, board_name: str, board_type: str) -> List[str]:
        """获取板块成份股"""
        try:
            if board_type == 'concept':
                df = ak.stock_board_concept_cons_em(symbol=board_name)
            elif board_type == 'industry':
                df = ak.stock_board_industry_cons_em(symbol=board_name)
            elif board_type == 'sector':
                # 地域板块暂时返回空列表，因为akshare没有对应的接口
                return []
            else:
                logger.warning(f"未知的板块类型: {board_type}")
                return []
            
            if df.empty:
                return []
            
            # 提取股票代码
            stock_codes = []
            if '代码' in df.columns:
                stock_codes = df['代码'].tolist()
            elif '股票代码' in df.columns:
                stock_codes = df['股票代码'].tolist()
            
            # 格式化股票代码（确保6位）
            formatted_codes = []
            for code in stock_codes:
                if isinstance(code, str) and code.isdigit():
                    formatted_codes.append(code.zfill(6))
                elif isinstance(code, (int, float)):
                    formatted_codes.append(str(int(code)).zfill(6))
            
            return formatted_codes
            
        except Exception as e:
            logger.error(f"获取板块 {board_name} 成份股失败: {e}")
            return []
    
    def sync_board_tags(self, board_df: pd.DataFrame) -> Dict[str, Any]:
        """同步板块标签"""
        result = {
            'success': False,
            'total_boards': len(board_df),
            'processed_boards': 0,
            'total_tags': 0,
            'total_stock_tags': 0,
            'error': None
        }
        
        try:
            logger.info(f"开始同步 {len(board_df)} 个板块标签...")
            
            processed_boards = 0
            total_tags = 0
            total_stock_tags = 0
            
            for _, board in board_df.iterrows():
                try:
                    board_name = board['name']
                    board_type = board['type']
                    
                    logger.info(f"处理板块: {board_name} ({board_type})")
                    
                    # 1. 插入或更新标签
                    tag_data = {
                        'name': board_name,
                        'type': board_type,
                        'source': 'akshare'
                    }
                    self.tag_dao.insert_or_update_tag(tag_data)
                    total_tags += 1
                    
                    # 2. 获取标签ID
                    tag_id = self.tag_dao.get_tag_id(board_name, board_type)
                    if not tag_id:
                        logger.error(f"无法获取标签ID: {board_name}")
                        continue
                    
                    # 3. 获取板块成份股
                    stock_codes = self.get_board_stocks(board_name, board_type)
                    
                    if stock_codes:
                        # 4. 批量插入股票标签关联
                        stock_tag_list = []
                        for stock_code in stock_codes:
                            stock_tag_list.append({
                                'code': stock_code,
                                'tag_id': tag_id,
                                'source': 'akshare'
                            })
                        
                        if stock_tag_list:
                            # 分批处理
                            for i in range(0, len(stock_tag_list), self.batch_size):
                                batch = stock_tag_list[i:i + self.batch_size]
                                synced_count = self.tag_dao.batch_insert_or_update_stock_tags(batch)
                                total_stock_tags += synced_count
                        
                        logger.info(f"✓ {board_name}: {len(stock_codes)} 只股票")
                    else:
                        logger.warning(f"⚠ {board_name}: 无成份股数据")
                    
                    processed_boards += 1
                    
                    # 添加延迟
                    if self.request_delay > 0:
                        import time
                        time.sleep(self.request_delay)
                        
                except Exception as e:
                    logger.error(f"处理板块 {board_name} 失败: {e}")
                    continue
            
            result['success'] = True
            result['processed_boards'] = processed_boards
            result['total_tags'] = total_tags
            result['total_stock_tags'] = total_stock_tags
            
            logger.info(f"板块标签同步完成！处理 {processed_boards}/{len(board_df)} 个板块，"
                       f"标签 {total_tags} 个，股票标签关联 {total_stock_tags} 个")
            
        except Exception as e:
            result['error'] = str(e)
            logger.error(f"板块标签同步失败: {e}")
        
        return result
    
    def sync_all_tags(self) -> Dict[str, Any]:
        """同步所有标签"""
        start_time = datetime.now()
        result = {
            'success': False,
            'concept_result': None,
            'industry_result': None,
            'sector_result': None,
            'duration': 0,
            'error': None
        }
        
        try:
            logger.info("开始同步所有股票标签...")
            
            # 同步概念板块
            logger.info("=" * 50)
            logger.info("同步概念板块标签")
            logger.info("=" * 50)
            concept_df = self.get_concept_boards()
            concept_result = self.sync_board_tags(concept_df)
            result['concept_result'] = concept_result
            
            # 同步行业板块
            logger.info("=" * 50)
            logger.info("同步行业板块标签")
            logger.info("=" * 50)
            industry_df = self.get_industry_boards()
            industry_result = self.sync_board_tags(industry_df)
            result['industry_result'] = industry_result
            
            # 同步地域板块
            logger.info("=" * 50)
            logger.info("同步地域板块标签")
            logger.info("=" * 50)
            sector_df = self.get_sector_boards()
            sector_result = self.sync_board_tags(sector_df)
            result['sector_result'] = sector_result
            
            # 检查所有结果
            all_success = (concept_result['success'] and 
                          industry_result['success'] and 
                          sector_result['success'])
            
            result['success'] = all_success
            
            if all_success:
                logger.info("=" * 60)
                logger.info("所有标签同步成功！")
                logger.info(f"概念板块: {concept_result['processed_boards']} 个")
                logger.info(f"行业板块: {industry_result['processed_boards']} 个")
                logger.info(f"地域板块: {sector_result['processed_boards']} 个")
                logger.info("=" * 60)
            else:
                logger.error("部分标签同步失败")
            
        except Exception as e:
            result['error'] = str(e)
            logger.error(f"标签同步失败: {e}")
        
        finally:
            end_time = datetime.now()
            result['duration'] = (end_time - start_time).total_seconds()
            logger.info(f"标签同步耗时: {result['duration']:.2f} 秒")
        
        return result
    
    def get_sync_summary(self) -> Dict[str, Any]:
        """获取同步摘要信息"""
        try:
            # 获取标签统计
            total_tags = self.db_manager.execute_query('SELECT COUNT(*) as count FROM tags')[0]['count']
            concept_tags = self.db_manager.execute_query("SELECT COUNT(*) as count FROM tags WHERE type = 'concept'")[0]['count']
            industry_tags = self.db_manager.execute_query("SELECT COUNT(*) as count FROM tags WHERE type = 'industry'")[0]['count']
            sector_tags = self.db_manager.execute_query("SELECT COUNT(*) as count FROM tags WHERE type = 'sector'")[0]['count']
            
            # 获取股票标签关联统计
            total_stock_tags = self.db_manager.execute_query('SELECT COUNT(*) as count FROM stock_tag')[0]['count']
            tagged_stocks = self.db_manager.execute_query('SELECT COUNT(DISTINCT code) as count FROM stock_tag')[0]['count']
            
            return {
                'total_tags': total_tags,
                'concept_tags': concept_tags,
                'industry_tags': industry_tags,
                'sector_tags': sector_tags,
                'total_stock_tags': total_stock_tags,
                'tagged_stocks': tagged_stocks,
                'last_sync': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }
            
        except Exception as e:
            logger.error(f"获取同步摘要失败: {e}")
            return {
                'error': str(e),
                'total_tags': 0,
                'concept_tags': 0,
                'industry_tags': 0,
                'sector_tags': 0,
                'total_stock_tags': 0,
                'tagged_stocks': 0
            }

def main():
    """主函数"""
    try:
        # 创建同步器
        sync = StockTagSync()
        
        # 执行同步
        result = sync.sync_all_tags()
        
        # 打印结果
        if result['success']:
            print("=" * 60)
            print("A股股票标签同步成功！")
            
            if result['concept_result']:
                print(f"概念板块: {result['concept_result']['processed_boards']} 个")
            if result['industry_result']:
                print(f"行业板块: {result['industry_result']['processed_boards']} 个")
            if result['sector_result']:
                print(f"地域板块: {result['sector_result']['processed_boards']} 个")
            
            print(f"同步耗时: {result['duration']:.2f} 秒")
            print("=" * 60)
        else:
            print("=" * 60)
            print("A股股票标签同步失败！")
            print(f"错误信息: {result['error']}")
            print("=" * 60)
        
        # 显示同步摘要
        summary = sync.get_sync_summary()
        if 'error' not in summary:
            print("\n数据库当前状态:")
            print(f"总标签数: {summary['total_tags']}")
            print(f"  概念标签: {summary['concept_tags']}")
            print(f"  行业标签: {summary['industry_tags']}")
            print(f"  地域标签: {summary['sector_tags']}")
            print(f"股票标签关联: {summary['total_stock_tags']}")
            print(f"已标签股票: {summary['tagged_stocks']}")
        
    except Exception as e:
        logger.error(f"程序执行失败: {e}")
        print(f"程序执行失败: {e}")

if __name__ == "__main__":
    main()
